import cv2
import numpy as np
# import torch

class __Colors:
    def __init__(self):
        hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB',
               '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb('#' + c) for c in hex]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))


get_color = __Colors()
class_names = ['blue', 'yellow', 'white', 'red', 'none']


def draw_bbox(imgs, results, draw_label=True, line_thickness=3):
    single = False
    if isinstance(imgs, np.ndarray):
        imgs = [imgs]
        single = True
    tf = max(line_thickness - 1, 1)
    for img, result in zip(imgs, [results]):
        # print(img.shape)
        if result is not None:
            # print(result.shape)
            flag = result.shape[-1] == 7
            for *xywh, conf, cls in result:
                c1 = (int(xywh[0]), int(xywh[1]))
                c2 = (int(xywh[2]), int(xywh[3]))
                color = get_color(int(cls), True)
                cv2.rectangle(img, c1, c2, color, line_thickness, cv2.LINE_AA)
                if draw_label:
                    label = f'{class_names[int(cls)]} {conf * (xywh[-1] if flag else 1):.2f}'
                    t_size = cv2.getTextSize(label, 0, fontScale=line_thickness / 3, thickness=tf)[0]
                    c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 3
                    cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA)  # filled
                    cv2.putText(img, label, (c1[0], c1[1] - 2), 0, line_thickness / 3, [225, 255, 255],
                                thickness=tf, lineType=cv2.LINE_AA)

    return imgs[0] if single else imgs


# if __name__ == '__main__':
    # result = torch.tensor([[2.63000e+02, 1.26000e+02, 3.10000e+02, 1.85000e+02, 8.40820e-01, 2.00000e+00],
    #     [1.86000e+02, 9.50000e+01, 2.36000e+02, 1.60000e+02, 8.38867e-01, 2.00000e+00],
    #     [8.27000e+02, 7.00000e+01, 8.72000e+02, 1.27000e+02, 8.38867e-01, 2.00000e+00],
    #     [1.00000e+01, 1.10000e+02, 6.00000e+01, 1.75000e+02, 8.32031e-01, 2.00000e+00],
    #     [4.29000e+02, 9.80000e+01, 4.77000e+02, 1.61000e+02, 8.30078e-01, 2.00000e+00],
    #     [7.65000e+02, 5.30000e+01, 8.22000e+02, 1.30000e+02, 8.25684e-01, 2.00000e+00],
    #     [9.31000e+02, 6.40000e+01, 1.00200e+03, 1.54000e+02, 8.25684e-01, 2.00000e+00],
    #     [6.92000e+02, 9.00000e+01, 7.43000e+02, 1.59000e+02, 8.20801e-01, 2.00000e+00],
    #     [9.50000e+01, 1.37000e+02, 1.42000e+02, 1.96000e+02, 8.06641e-01, 2.00000e+00],
    #     [5.88000e+02, 1.14000e+02, 6.38000e+02, 1.74000e+02, 8.01758e-01, 2.00000e+00],
    #     [8.61000e+02, 8.50000e+01, 9.18000e+02, 1.56000e+02, 7.96875e-01, 2.00000e+00],
    #     [3.35000e+02, 8.80000e+01, 3.82000e+02, 1.52000e+02, 7.85156e-01, 2.00000e+00],
    #     [5.03000e+02, 9.60000e+01, 5.52000e+02, 1.60000e+02, 7.82715e-01, 2.00000e+00]], device='cuda:0')
    # img = cv2.imread("asset/00000.jpg")
    # img = draw_bbox(imgs=img, results=result)
    # cv2.imshow("1",img)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()